package edu.cornell.cs4740.postag.hmm;

import java.util.ArrayList;
import java.util.List;

import edu.cornell.cs4740.postag.parsing.Sentence;
import edu.cornell.cs4740.postag.utils.Utilities;

//assumes using bigrams
public class HMM_viterbi {

	// [i][j] is the probability of word i given p-o-s j
	private double[][] word_pos_probs;
	// [i][j] is the probability of pos i given pos j immediately before it in
	// sequence
	private double[][] bigram_probs;
	// [i][j][k] is the probability of pos i given pos k,j immediately before it
	// in sequence
	private double[][][] trigram_probs;
	// [i] is the probability of pos i being the first word of a sentence
	private double[] first_pos_probs;
	// [i][j] is the probability of pos i being immediately after pos j at the
	// beginning of the sentence
	private double[][] second_pos_probs;
	// s=sentence to be labeled with p-o-s
	private Sentence s;
	// c = num possible p-o-s'
	private int c;

	// score[i][j] is the highest possible probability of word i in s to be
	// p-o-s j in pos_lexicon
	double[][] score;
	// bptr[i][j] is the index of the p-o-s in pos_lexicon of the word that
	// comes immediately before word i with p-o-s j in ij's most probable path
	// index before first word is -1
	int[][] bptr;
	// 2 before word i with pos j in mpp
	int[][] bptr_2;

	public HMM_viterbi(int c, double[][] word_given_pos_probs,
			double[][] bigram_probabilities,
			double[][][] trigram_probabilities, double[] first_pos_probs,
			double[][] second_pos_probs) {
		this.bigram_probs = bigram_probabilities;
		this.trigram_probs = trigram_probabilities;
		this.word_pos_probs = word_given_pos_probs;
		this.first_pos_probs = first_pos_probs;
		this.second_pos_probs = second_pos_probs;
		this.c = c;
	}

	public HMM_viterbi(int c, double[][] word_given_pos_probs,
			double[][] bigram_probabilities,
			double[][][] trigram_probabilities, double[] first_pos_probs,
			double[][] second_pos_probs, Sentence s) {
		this(c, word_given_pos_probs, bigram_probabilities,
				trigram_probabilities, first_pos_probs, second_pos_probs);
		this.s = s;
	}

	public void setSentence(Sentence s) {
		this.s = s;
	}

	private void init(boolean useTrigrams) {
		score = new double[s.getLength()][c];
		bptr = new int[s.getLength()][c];
		bptr_2 = new int[s.getLength()][c];
		double tempScore;

		for (int i = 0; i < c; i++) {
			score[0][i] = first_pos_probs[i] * getWordPosProb(0, i);
			bptr[0][i] = -1;
			bptr_2[0][i] = -1;
		}
		if (useTrigrams) {
			for (int i = 0; i < c; i++) {
				for (int j = 0; j < c; j++) {
					if (i == 0) {
						score[1][j] = second_pos_probs[j][i]
								* getWordPosProb(0, i) * getWordPosProb(1, j);
						bptr[1][j] = i;
						bptr_2[1][j] = -1;
					} else {
						tempScore = second_pos_probs[j][i]
								* getWordPosProb(0, i) * getWordPosProb(1, j);
						if (tempScore > score[1][j]) {
							score[1][j] = tempScore;
							bptr[1][j] = i;
						}
					}
				}
			}
		}
	}

	private double getWordPosProb(int i, int t) {
		int p = Utilities.word_to_index(s.getWord(i).getWord());
		if (p == -1)
			return 1.0 / c; // is 1 over c correct? does it even matter?
		return word_pos_probs[p][t];
	}

	// true = trigrams
	// false = bigrams
	public List<String> findMostProbablePOSs(boolean useTrigrams) {
		double maxScore, tempScore;
		int maxIndex, maxIndexJ, maxIndexI;

		useTrigrams = useTrigrams && s.getLength()>1;
		
		init(useTrigrams);

		// build probability table and paths
		if (!useTrigrams) {
			for (int t = 1; t < s.getLength(); t++) {
				for (int i = 0; i < c; i++) {
					maxScore = Double.NEGATIVE_INFINITY;
					maxIndex = -1;
					for (int j = 0; j < c; j++) {
						tempScore = score[t - 1][j] * bigram_probs[i][j]
								* getWordPosProb(t, i);
						if (tempScore > maxScore) {
							maxScore = tempScore;
							maxIndex = j;
						}
					}
					score[t][i] = maxScore;
					bptr[t][i] = maxIndex;
				}
			}
		} else {
			for (int t = 2; t < s.getLength(); t++) {
				for (int k = 0; k < c; k++) {
					maxScore = Double.NEGATIVE_INFINITY;
					maxIndexI = -1;
					maxIndexJ = -1;
					for (int i = 0; i < c; i++) {
						for (int j = 0; j < c; j++) {
							tempScore = score[t - 2][i]
									* trigram_probs[k][i][j]
									* getWordPosProb(t - 1, j)
									* getWordPosProb(t, k);
							if (tempScore > maxScore) {
								maxScore = tempScore;
								maxIndexJ = j;
								maxIndexI = i;
							}
						}
					}
					score[t][k] = maxScore;
					bptr[t][k] = maxIndexJ;
					bptr_2[t][k] = maxIndexI;
				}
			}
		}

		// identify most probable sequence
		int[] i_mpp = new int[s.getLength()];
		maxScore = Double.NEGATIVE_INFINITY;
		maxIndex = -1;
		for (int i = 0; i < c; i++) {
			tempScore = score[s.getLength() - 1][i];
			if (tempScore > maxScore) {
				maxScore = tempScore;
				maxIndex = i;
			}
		}
		i_mpp[i_mpp.length - 1] = maxIndex; // end of most probable path

		// backtrack from there
		for (int i = s.getLength() - 2; i >= 0; i--) {
			i_mpp[i] = bptr[i + 1][i_mpp[i + 1]];
			if (useTrigrams && i >= 1) {
				i_mpp[i - 1] = bptr_2[i + 1][i_mpp[i + 1]];
				i--;
			}
		}

		// convert indexes to actual p-o-s
		List<String> mpp = new ArrayList<String>();
		for (int i = 0; i < i_mpp.length; i++) {
			mpp.add(Utilities.index_to_pos(i_mpp[i]));
		}

		return mpp;
	}
}
